import numpy as np
import matplotlib.pyplot as plt

np.random.seed(1)


def nms(boxes, thresh=0.7):
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]
    P = boxes[:, 4]
    A = (x2 - x1 + 1) * (y2 - y1 + 1)

    order = P.argsort()[::-1]

    # idx_keep_arr = np.array([])  # IndexError: arrays used as indices must be of integer (or boolean) type
    idx_keep_arr = np.array([], dtype=np.int32)
    idx_uncertain_arr = order.copy()
    while True:
        if len(idx_uncertain_arr) == 0:
            break
        idx_keep = idx_uncertain_arr[0]
        idx_keep_arr = np.append(idx_keep_arr, [idx_keep])
        idx_uncertain_arr = idx_uncertain_arr[1:]
        print('Before')
        print('idx_keep_arr', idx_keep_arr)
        print('idx_uncertain_arr', idx_uncertain_arr)
        if len(idx_uncertain_arr) == 0:
            break

        max_x1 = np.maximum(x1[idx_keep], x1[idx_uncertain_arr])
        min_x2 = np.minimum(x2[idx_keep], x2[idx_uncertain_arr])
        max_y1 = np.maximum(y1[idx_keep], y1[idx_uncertain_arr])
        min_y2 = np.minimum(y2[idx_keep], y2[idx_uncertain_arr])
        xx = min_x2 - max_x1 + 1
        xx = np.maximum(xx, 0)
        yy = min_y2 - max_y1 + 1
        yy = np.maximum(yy, 0)
        I = xx * yy
        U = A[idx_keep] + A[idx_uncertain_arr] - I
        iou = I / (U + 1e-10)
        idx_uncertain_arr = idx_uncertain_arr[iou < thresh]

        print('After')
        print('idx_keep_arr', idx_keep_arr)
        print('idx_uncertain_arr', idx_uncertain_arr)
    return idx_keep_arr


def draw_boxes(boxes, idx):
    global spn
    spn += 1
    plt.subplot(spr, spc, spn)
    x1 = boxes[idx, 0]
    y1 = boxes[idx, 1]
    x2 = boxes[idx, 2]
    y2 = boxes[idx, 3]
    plt.plot([x1, x2], [y1, y1], color='k')
    plt.plot([x1, x2], [y2, y2], color='k')
    plt.plot([x1, x1], [y1, y2], color='k')
    plt.plot([x2, x2], [y1, y2], color='k')
    for i in idx:
        plt.annotate(f'#{i + 1}', xy=[boxes[i, 0], boxes[i, 1]])
    plt.ylim(plt.ylim()[::-1])


spr = 1
spc = 2
spn = 0

if '__main__' == __name__:
    boxes = np.array([
        [100, 100, 210, 210, 0.72],  # 0
        [280, 290, 420, 420, 0.8],  # 1
        [220, 220, 320, 330, 0.92],  # 2
        [105, 90, 220, 210, 0.71],  # 3
        [230, 240, 325, 330, 0.81],  # 4
        [305, 300, 420, 420, 0.9],  # 5
        [215, 225, 305, 328, 0.6],  # 6
        [150, 260, 290, 400, 0.99],  # 7
        [102, 108, 208, 208, 0.72]])  # 8  #9个框
    left_idx = nms(boxes)
    plt.figure(figsize=[12, 6])
    draw_boxes(boxes, list(range(len(boxes))))
    draw_boxes(boxes, left_idx)
    plt.show()
